package com.example.demo.sharding;

import com.mybatisflex.core.datasource.DataSourceShardingStrategy;
import org.springframework.util.StringUtils;

import java.lang.reflect.Method;

/**
 * 多数据源分片策略
 */
public class DBShardingStrategy implements DataSourceShardingStrategy {

    public static ThreadLocal<String> dbLocal = new ThreadLocal<>();

    @Override
    public String doSharding(String currentDataSourceKey, Object mapper, Method mapperMethod, Object[] methodArgs) {
        String db = dbLocal.get();
        if(db == null){
            // 如果 mapper 的方法属于 增删改，使用 master 数据源
            if (StringUtils.startsWithIgnoreCase(mapperMethod.getName(),"insert") ||
                StringUtils.startsWithIgnoreCase(mapperMethod.getName(),"delete") ||
                StringUtils.startsWithIgnoreCase(mapperMethod.getName(),"update") ){

                return "master";
            }
        }
        dbLocal.remove();
        return db;
    }
}
